package halftone.errdiff.pipeline

import chisel3._
import chisel3.util.{Counter, Decoupled}
import halftone.ErrDiffConfig
import halftone.errdiff.ELUT
import halftone.errdiff.ErrDiffCorePipParam._

class ThreshCalcPip(config: ErrDiffConfig) extends Module {
  val io = IO(new Bundle {
    val in  = Flipped(Decoupled(new ErrorIn2ThreshCalcPip(config.pixelWidth, config.posWidth, config.errorWidth)))
    val out = Decoupled(new ThreshCalc2ErrorOutPip(config.posWidth, config.errorWidth))
    // Forwarding
    val err   = Input(SInt(config.errorWidth.W))
    val err_r = Output(Vec(2, SInt(config.errorWidth.W)))
    val err_fc = Input(SInt(config.errorWidth.W))
    val stop = Input(Bool())
  })

  // Registers(for value storage and state presentation)
  val pos   = Reg(UInt(config.posWidth.W))
  val pix   = Reg(UInt((config.pixelWidth + 2).W)) // expand 2 bits
  val err   = Reg(SInt((config.errorWidth + 2).W)) // expand 2 bits
  val errIn = Reg(Vec(4, SInt(config.errorWidth.W)))

  val busy        = RegInit(false.B)
  val resultValid = RegInit(false.B)

  io.out.valid := resultValid
  io.in.ready  := true.B

  /*
   * Lookup Table
   */
  val elut = Module(new ELUT(config.errorWidth))

  /*
   * Compare with Threshold
   */
  val realPix = pix.asSInt + err
  val binOut  = Mux(realPix < config.threshold.S, 0.U, 1.U)

  /*
   * Calculate errors to be diffused
   */
  elut.io.err := realPix - Mux(binOut.asBool, 255.S(config.pixelWidth.W), 0.S)

  /*
   * Forwarding
   */
  /**
   * Get current error from EO by Forwarding,
   * And special dealing with first column error
   */
  err := Mux(config.isFirstColumn(pos), io.err_fc, io.err)
  /** Send read errors to EI by Forwarding */
  io.err_r(0) := errIn(1) + elut.io.out(1) // belowRight to below
  io.err_r(1) := errIn(2) + elut.io.out(2) // below to belowLeft

  // Emit outputs
  io.out.bits.pos    := pos
  io.out.bits.bval   := binOut
  io.out.bits.errOut := (errIn.zip(elut.io.out)).map { case (eIn, eOut) => eIn + eOut }

  val (cnt, cntWrap) = Counter(busy && !resultValid, CycleNr)

  when(cntWrap && !io.stop) { resultValid := true.B }
  when(io.out.fire) {
    resultValid := false.B
  }
  when(io.in.fire) {
    val inBundle = io.in.deq()
    pos   := inBundle.pos
    pix   := inBundle.pix
    errIn := inBundle.errIn
    busy  := true.B
  }
}
